
from PyQt4.QtCore import *
from PyQt4.QtGui import *

from qgis.core import *
from qgis.gui import *

from math import pi,sin,cos,sqrt

class SelectAction(QAction):
    def __init__(self,iface,menuTitle,whatsThis):
        mw = iface.mainWindow()
        QAction.__init__(self,menuTitle,mw)
        self.setWhatsThis(whatsThis)
        self.iface=iface
        QObject.connect(self, SIGNAL("activated()"), self.doit)

        QObject.connect(iface,SIGNAL("currentLayerChanged(QgsMapLayer *)"),self.setForLayer)
        # make sure we are properly initialised for the current layer:
        self.setForLayer(iface.mapCanvas().currentLayer())
        return None

    def setForLayer(self,layer):
        if not layer:
            self.setEnabled(False)
            return
        if layer.isValid():
            if layer.type() == layer.VectorLayer:
                self.setEnabled(True)
            else:
                self.setEnabled(False)
        else:
            self.setEnabled(False)
        return None

class InvertAction(SelectAction):
    def __init__(self,iface):
        SelectAction.__init__(self,iface,"Invert Selection","Invert the selection set")
        return None
    def doit(self):
        layer = self.iface.mapCanvas().currentLayer()
        howMany = layer.featureCount()
        current = layer.selectedFeaturesIds()
        new = list(set(range(howMany))-set(current))
        layer.setSelectedFeatures(new)
        return None

class ClearAction(SelectAction):
    def __init__(self,iface):
        SelectAction.__init__(self,iface,"Select None","Clear the selection set")
        return None
    def doit(self):
        layer = self.iface.mapCanvas().currentLayer()
        layer.setSelectedFeatures([])
        return None
    
class AllAction(SelectAction):
    def __init__(self,iface):
        SelectAction.__init__(self,iface,"Select All","Select all features")
        return None
    def doit(self):
        layer = self.iface.mapCanvas().currentLayer()
        layer.setSelectedFeatures(range(layer.featureCount()))
        return None

class PolyAction(SelectAction):
    def __init__(self,iface):
        SelectAction.__init__(self,iface,"Select polygon","Select features by drawing a polygon")
        return None
    def doit(self):
        canvas = self.iface.mapCanvas()
        self.tool = PolygonTool(canvas,self.iface)
        canvas.setMapTool(self.tool)
        return None

class FreehandAction(SelectAction):
    def __init__(self,iface):
        SelectAction.__init__(self,iface,"Select by freehand","Select features by freehand")
        return None
    def doit(self):
        canvas = self.iface.mapCanvas()
        self.tool = FreehandTool(self.iface)
        canvas.setMapTool(self.tool)
        return None

class CircleAction(SelectAction):
    def __init__(self,iface):
        SelectAction.__init__(self,iface,"Select by radius","Select features by center and edge")
        return None
    def doit(self):
        canvas = self.iface.mapCanvas()
        self.tool = CircleTool(self.iface)
        canvas.setMapTool(self.tool)
        return None

class Circle2Action(SelectAction):
    def __init__(self,iface):
        SelectAction.__init__(self,iface,"Select by radius","Select features by center and edge")
        return None
    def doit(self):
        canvas = self.iface.mapCanvas()
        self.tool = CircleTool2(self.iface)
        canvas.setMapTool(self.tool)
        return None

class CircleTool(QgsMapTool):
    def __init__(self,iface):
        canvas = iface.mapCanvas()
        QgsMapTool.__init__(self,canvas)
        self.canvas = canvas
        self.iface = iface
        self.status=0
        return None
    def canvasPressEvent(self,e):
        if not e.button() == Qt.LeftButton:
            return
        sb = self.iface.mainWindow().statusBar()
        sb.showMessage("Click center, click edge - ctrl to add, shift to substract")
        layer = self.canvas.currentLayer()
        self.status = 1
        self.center = self.toLayerCoordinates(layer,e.pos())
        self.rb=QgsRubberBand(self.canvas,True)
        return
    
    def canvasMoveEvent(self,e):
        if not self.status == 1:
            return
        # construct a circle with N segments
        layer = self.canvas.currentLayer()
        cp = self.toLayerCoordinates(layer,e.pos())
        rbcircle(self.rb,self.center,cp,40)
        r = sqrt(self.center.sqrDist(cp))
        sb = self.iface.mainWindow().statusBar()
        sb.showMessage("Center: %s %s, radius: %s" % (str(self.center.x()),str(self.center.y()),str(r)))

    def canvasReleaseEvent(self,e):
        if not e.button() == Qt.LeftButton:
            return
        layer = self.canvas.currentLayer()
        self.edge = self.toLayerCoordinates(layer,e.pos())
        selectRubberBand(layer, self.rb, e.modifiers())
        self.rb.reset()
        self.status=0
        

def rbcircle(rb,center,edgePoint,N):
    r = sqrt(center.sqrDist(edgePoint))
    rb.reset()
    for itheta in range(N+1):
        theta = itheta*(2.0 * pi/N)
        rb.addPoint(QgsPoint(center.x()+r*cos(theta),center.y()+r*sin(theta)))
    return 

class FreehandTool(QgsMapTool):
    def __init__(self,iface):
        canvas=iface.mapCanvas()
        QgsMapTool.__init__(self,canvas)
        # these dont get initted until first mouse-down
        self.rb=None
        self.layer=None
        self.canvas=canvas
        self.iface=iface
        #debugging self.events=[]
        return None
    def canvasPressEvent(self,e):
        #if not e.button()==Qt.LeftButton:
        #return
        # check this because sometimes we get two presses without a release
        #debugging self.events.append(('press',e.type()))
        if self.rb:
            return
        sb = self.iface.mainWindow().statusBar()
        sb.showMessage("Draw polygon, hold ctrl to add, shift to subtract")
        self.layer = self.canvas.currentLayer()
        self.rb=QgsRubberBand(self.canvas,True)
        p=self.toLayerCoordinates(self.layer,e.pos())
        self.rb.addPoint(p)
        return 
    def canvasMoveEvent(self,e):
        #debugging self.events.append(('move',e.type()))
        if not self.rb:
            return
        p=self.toLayerCoordinates(self.layer,e.pos())
        self.rb.addPoint(p)
        return
    def canvasReleaseEvent(self,e):
        #debugging self.events.append(('release',e.type()))
        if not self.rb:
            return
        #if e.button() == Qt.LeftButton:
        if self.rb.numberOfVertices() > 2:
            selectRubberBand(self.layer,self.rb,e.modifiers())
        self.rb.reset()
        self.rb=None
        return

    #debugging def deactivate(self):
    #eventList(self.events)
    #return

class PolygonTool(QgsMapTool):
    def __init__(self,canvas,iface):
        QgsMapTool.__init__(self,canvas)
        self.iface=iface
        self.canvas=canvas
        self.rb = QgsRubberBand(canvas, True)
        self.layer = canvas.currentLayer()
        return None

    def deactivate(self):
        self.rb.reset()
        return None
    
    def canvasPressEvent(self,e):
        #print "Got press event: ",str(e)
        self.layer = self.canvas.currentLayer()
        sb = self.iface.mainWindow().statusBar()
        sb.showMessage("Left-click points, right-click to finish, ctrl to add, shift to subtract")
        if e.button() == Qt.LeftButton:
            p=self.toLayerCoordinates(self.layer,e.pos())
            self.rb.addPoint(p)
        else:
            if self.rb.numberOfVertices() > 2:
                selectRubberBand(self.layer,self.rb,e.modifiers())
                # clear the rubber band
                self.rb.reset()
        return None

    def canvasMoveEvent(self,e):
        if self.rb.numberOfVertices() > 0:
            self.rb.removeLastPoint(0)
            self.layer = self.canvas.currentLayer()
            p = self.toLayerCoordinates(self.layer,e.pos())
            self.rb.addPoint(p)
        return None

#     def canvasReleaseEvent(self,e):
#         #print "Got release event: ",str(e)
#         return None
#     def keyPressEvent(self,e):
#         #print "Got key press event: ",str(e)
#         return None


def eventList(estuples):
    "used in debugging"
    m=QgsMessageOutput.createMessageOutput()
    m.setTitle("events")
    m.setMessage("Here:",QgsMessageOutput.MessageText)
    for e in estuples:
        m.appendMessage(e[0]+":"+str(e[1]))
    m.showMessage()
    return

def errorLayer(layer,msg):
    m=QgsMessageOutput.createMessageOutput()
    m.setTitle("Invalid layer")
    m.setMessage("The current layer is not a valid vector layer.", QgsMessageOutput.MessageText)
    m.appendMessage(msg)
    m.showMessage()
    return 

def selectRubberBand(layer, rb, modifiers):
    if not layer:
        errorLayer(layer,"Layer is None")
        return
    if not layer.type() == layer.VectorLayer:
        errorLayer(layer, "Layer is not a Vector Layer")
        return
    coords = []
    [coords.append(rb.getPoint(0,i)) for i in range(rb.numberOfVertices())]
    g=QgsGeometry().fromPolygon([coords])
    # now look over current layer features f and test f.geometry().intersects(g)
    provider = layer.dataProvider()
    feat = QgsFeature()
    selectedSet = []
    while provider.nextFeature(feat):
        geom = feat.geometry()
        if geom.intersects(g):
            selectedSet.append(feat.id())
    # maybe do a bounding box check first
    # control key is add to selection
    # shift key is remove from selection
    # no modifier is change selection
    if modifiers & Qt.ControlModifier:
        selectedSet = list(set(layer.selectedFeaturesIds()).union(selectedSet))
    elif modifiers & Qt.ShiftModifier:
        selectedSet = list(set(layer.selectedFeaturesIds()).difference(selectedSet))
    # update selected features
    layer.setSelectedFeatures(selectedSet)


class CircleTool2(QgsMapTool):
    def __init__(self,iface):
        canvas = iface.mapCanvas()
        QgsMapTool.__init__(self,canvas)
        self.canvas = canvas
        self.iface = iface
        self.status=0
        return None
    def canvasPressEvent(self,e):
        if not e.button() == Qt.LeftButton:
            return
        sb = self.iface.mainWindow().statusBar()
        sb.showMessage("Click center, click edge - ctrl to add, shift to substract")
        layer = self.canvas.currentLayer()
        self.status = 1
        transform = self.canvas.getCoordinateTransform()
        self.center = transform.toMapCoordinates(e.pos().x(),e.pos().y())
        self.rb=QgsRubberBand(self.canvas,True)
        return
    
    def canvasMoveEvent(self,e):
        if not self.status == 1:
            return
        # construct a circle with N segments
        layer = self.canvas.currentLayer()
        transform = self.canvas.getCoordinateTransform()
        cp = transform.toMapCoordinates(e.pos().x(),e.pos().y())
        rbcircle(self.rb,self.center,cp,40)
        r = sqrt(self.center.sqrDist(cp))
        sb = self.iface.mainWindow().statusBar()
        sb.showMessage("Center: %s %s, radius: %s" % (str(self.center.x()),str(self.center.y()),str(r)))

    def canvasReleaseEvent(self,e):
        if not e.button() == Qt.LeftButton:
            return
        layer = self.canvas.currentLayer()
        selectRubberBand(layer, self.rb, e.modifiers())
        self.rb.reset()
        self.status=0
        
